import weakref
from typing import Callable, Iterable, Iterator, Optional, Union

import matplotlib.pyplot as plt
import numpy as np
from IPython.display import display
from ipywidgets import widgets

from .live_channel import InputType, IntervalLiveChannel, IntervalLiveProducer, OutputType
from .live_figure import LiveFigure


class HistBufferChannel(IntervalLiveChannel):
    def __init__(self, *, interval: float = 1 / 30, bins: int = 10):
        super().__init__(interval=interval)
        self._bins = bins

        # buffers
        self._ys_logs_buffer = bytearray()
        self._ys_n = None

    def clear(self):
        self.send(None)

    def _update_buffer(self, update: InputType):
        if update is None:  # clear
            self._ys_logs_buffer.clear()
            self._ys_n = None
            return

        ys = update
        ys = np.asarray(ys, dtype=np.float32).reshape([-1])
        ys_n = np.shape(ys)[0]

        if self._ys_n is None:
            self._ys_n = ys_n
        elif self._ys_n != ys_n:
            raise ValueError(f"Expected ys_n={self._ys_n}, got {ys_n}")

        self._ys_logs_buffer.extend(ys.tobytes())

    def _read_buffer(self) -> OutputType:
        if len(self._ys_logs_buffer) == 0 or self._ys_n is None:
            return None  # clear
        ys_logs = np.frombuffer(self._ys_logs_buffer, dtype=np.float32).reshape([-1, self._ys_n])
        histograms = [np.histogram(ys_logs[:, yi], bins=self._bins) for yi in range(self._ys_n)]
        return histograms


class HistFigure(LiveFigure):
    def __init__(self, figure: Optional[plt.Figure] = None):
        super().__init__(figure)
        self.figure.set_size_inches(8, 4)
        self._axes = self.figure.gca()
        self._axes.grid()

        self._ys_bars: Optional[list[Iterable[plt.Rectangle]]] = None

    @staticmethod
    def _init_patches(ys_n, bars_n, axes: plt.Axes):
        return [axes.bar(
            np.zeros([bars_n]), np.zeros([bars_n]),
            edgecolor="black", linewidth=1,
            zorder=2 + (ys_n - yi) / ys_n
        ) for yi in range(ys_n)]

    def _update_gui(self, data: OutputType):
        if data is None:  # clear
            self._axes.cla()
            self._axes.grid()
            self._ys_bars = None
            return

        ys_histogram = data
        if self._ys_bars is None:
            ys_n = len(ys_histogram)
            if ys_n == 0:
                return
            bars_n = len(ys_histogram[0][0])
            if bars_n == 0:
                return
            self._ys_bars = self._init_patches(ys_n, bars_n, self._axes)

        max_height = 0
        min_x, max_x = None, None
        for y_bars, (y_heights, y_edges) in zip(self._ys_bars, ys_histogram):
            for i, bar in enumerate(y_bars):
                bar.set_x(y_edges[i])
                bar.set_width(y_edges[i + 1] - y_edges[i])
                bar.set_height(y_heights[i])
            max_height = max(max_height, np.max(y_heights))
            min_x = min(min_x, min(y_edges)) if min_x is not None else min(y_edges)
            max_x = max(max_x, max(y_edges)) if max_x is not None else max(y_edges)

        padding = (max_x - min_x) / 32
        self._axes.set_xlim(min_x - padding, max_x + padding)
        self._axes.set_ylim(0, max_height * 1.1)


class LiveHistogramWidgets:
    def __init__(self,
            producer: Union[Iterator[Iterable[float]], Callable[[], Iterable[float]]], *,
            produce_interval: float = 1 / 128,
            frame_interval: float = 1 / 32,

            bins: int = 20,

            figure: Optional[plt.Figure] = None,
            start: bool = True
    ):
        # ui widgets

        play_button = widgets.Button()
        play_button.description = "start"
        display(play_button)

        clear_button = widgets.Button()
        clear_button.description = "clear"
        display(clear_button)

        live_figure = HistFigure(
            figure=figure)

        # buffer

        live_buffer = HistBufferChannel(
            interval=frame_interval,
            bins=bins)
        live_buffer.subscribe(live_figure)

        # producer

        if hasattr(producer, '__next__'):
            prod_func = lambda: next(producer)
        elif callable(producer):
            prod_func = producer
        else:
            raise TypeError(f"The producer is expected to be an iterator (generator) or callable, got {producer}")

        live_producer = IntervalLiveProducer(prod_func,
            interval=produce_interval)
        live_producer.subscribe(live_buffer)

        # use weak ref to avoid hanging by jupyter as displaying widgets

        producer_ref = weakref.ref(live_producer)
        buffer_ref = weakref.ref(live_buffer)

        # play button

        def on_play_button_clicked(_):
            producer = producer_ref()
            if producer is None:
                return
            play_button.disabled = True
            if producer.alive:
                producer.stop()
                play_button.description = "resume"
            else:
                producer.start()
                play_button.description = "pause"
            play_button.disabled = False

        play_button.on_click(on_play_button_clicked)

        # clear button

        def on_clear_button_clicked(_):
            buffer = buffer_ref()
            if buffer is None:
                return
            buffer.clear()

        clear_button.on_click(on_clear_button_clicked)

        # start
        if start:
            play_button.click()

        # properties
        self.play_button = play_button
        self.clear_button = clear_button
        self.producer = live_producer
        self.buffer = live_buffer
        self.figure = live_figure
